import argparse
import logging
import os

import nibabel as nib
import numpy as np
import vtk
from vtkmodules.vtkCommonCore import (
    VTK_VERSION_NUMBER,
    vtkVersion
)
from vtkmodules.vtkCommonDataModel import (
    vtkDataObject,
    vtkDataSetAttributes
)
from vtkmodules.vtkFiltersCore import (
    vtkMaskFields,
    vtkThreshold,
    vtkWindowedSincPolyDataFilter
)
from vtkmodules.vtkFiltersGeneral import (
    vtkDiscreteFlyingEdges3D,
    vtkDiscreteMarchingCubes
)
from vtkmodules.vtkFiltersGeometry import vtkGeometryFilter
from vtkmodules.vtkIOGeometry import (
    vtkOBJWriter
)


def parse_args():
    parser = argparse.ArgumentParser(description="generate models from labels")

    parser.add_argument('--file_path', type=str, default=None, required=True, help='path to the file to be generate')
    parser.add_argument('--start_label', type=int, default=1,
                        help='start label index, if none we will use 1 as default')
    parser.add_argument('--end_label', type=int, default=None, required=False,
                        help='end label index, if none we will use the max label in the file')
    parser.add_argument('--save_root', type=str, default=None, required=True, help='path to save the generated models')
    parser.add_argument('--file_format', type=str, default='nii.gz',
                        help='file format, one of nii.gz, nii, mhd, raw, dcm')
    parser.add_argument('--model_format', type=str, default='obj',
                        help='model format, one of obj, now we only support .obj')
    parser.add_argument('--origin', type=int, default=None, help='origin of the dataset')
    parser.add_argument('--spacing', type=int, default=None, help='spacing of the dataset')
    parser.add_argument('--smoothing_iterations', type=int, default=15, help='smoothing iterations')
    parser.add_argument('--pass_band', type=float, default=0.001, help='pass band')
    parser.add_argument('--feature_angle', type=float, default=120.0, help='feature angle')
    parser.add_argument('--file_prefix', type=str, default='', help='file prefix')
    parser.add_argument('--class_names', type=str, nargs='*', default=None, help='class name')
    parser.add_argument('--remove_class', type=str, nargs='*', default=None, help='remove class')

    return parser.parse_args()


def make_reader(file_format='nii.gz'):
    if file_format == 'nii.gz':
        reader = vtk.vtkNIFTIImageReader()
    elif file_format == 'nii':
        reader = vtk.vtkNIFTIImageReader()
    elif file_format == 'mhd':
        reader = vtk.vtkMetaImageReader()
    elif file_format == 'raw':
        reader = vtk.vtkImageReader2()
    elif file_format == 'dcm':
        reader = vtk.vtkDICOMImageReader()
    else:
        print("Error: file_format must be one of 'nii.gz', 'nii', 'mhd', 'raw', 'dcm'")
        return
    return reader


def vtk_version_ok(major, minor, build):
    """
    Check the VTK version.

    :param major: Major version.
    :param minor: Minor version.
    :param build: Build version.
    :return: True if the requested VTK version is greater or equal to the actual VTK version.
    """
    needed_version = 10000000000 * int(major) + 100000000 * int(minor) + int(build)
    try:
        vtk_version_number = VTK_VERSION_NUMBER
    except AttributeError:  # as error:
        ver = vtkVersion()
        vtk_version_number = 10000000000 * ver.GetVTKMajorVersion() + 100000000 * ver.GetVTKMinorVersion() \
                             + ver.GetVTKBuildVersion()
    if vtk_version_number >= needed_version:
        return True
    else:
        return False


def cheak_list_file_format_fit(file_path_list):
    formats = []
    for i in file_path_list:
        formats.append(os.path.basename(i).split('.')[-1])
    return not len(set(formats)) > 1


def generate_one(file_name, start_label, end_label, save_root, file_format='nii.gz',
                 model_format='obj', Origin=(0, 0, 0), Spacing=(1, 1, 1),
                 smoothing_iterations=15, pass_band=0.001, feature_angle=120.0,
                 file_prefix='', class_names=None, remove_class=None):
    assert file_format in ['nii.gz', 'nii', 'mhd', 'raw',
                           'dcm'], "file_format must be one of 'nii.gz', 'nii', 'mhd', 'raw', 'dcm'"
    # TODO add other model formats
    assert model_format in ['obj'], "model_format must be one of 'obj'"

    if file_prefix == '':
        file_prefix = os.path.basename(file_name).split('.')[0]

    if os.path.exists(save_root) is False:
        os.makedirs(save_root)

    # **************** INIT ****************

    use_flying_edges = vtk_version_ok(8, 2, 0)

    if start_label > end_label:
        logging.error("start_label must be less than end_label")
        return

    reader = make_reader(file_format)
    histogram = vtk.vtkImageAccumulate()

    if use_flying_edges:
        try:
            using_marching_cubes = False
            # vtkDiscreteFlyingEdges3D 从 3D 图像数据（体积）生成等值面
            discrete_cubes = vtkDiscreteFlyingEdges3D()
        except AttributeError:
            using_marching_cubes = True
            discrete_cubes = vtkDiscreteMarchingCubes()
    else:
        using_marching_cubes = True
        discrete_cubes = vtkDiscreteMarchingCubes()

    smoother = vtkWindowedSincPolyDataFilter()
    selector = vtkThreshold()
    scalars_off = vtkMaskFields()

    geometry = vtkGeometryFilter()

    writer = vtkOBJWriter()

    # 1) Read the meta file
    reader.SetFileName(file_name)
    # 2) Generate a histogram of the labels
    # TODO 直方图的参数
    histogram.SetInputConnection(reader.GetOutputPort())
    histogram.SetComponentExtent(0, end_label, 0, 0, 0, 0)
    histogram.SetComponentOrigin(Origin[0], Origin[1], Origin[2])
    histogram.SetComponentSpacing(Spacing[0], Spacing[1], Spacing[2])
    histogram.Update()

    # 3) Generate models from the labeled volume
    discrete_cubes.SetInputConnection(reader.GetOutputPort())
    discrete_cubes.GenerateValues(end_label - start_label + 1, start_label, end_label)

    # 4) Smooth the models
    smoother.SetInputConnection(discrete_cubes.GetOutputPort())
    smoother.SetNumberOfIterations(smoothing_iterations)
    smoother.BoundarySmoothingOff()
    smoother.FeatureEdgeSmoothingOff()
    smoother.SetFeatureAngle(feature_angle)
    smoother.SetPassBand(pass_band)
    smoother.NonManifoldSmoothingOn()
    smoother.NormalizeCoordinatesOn()
    smoother.Update()

    # 5) Output each model into a separate file
    selector.SetInputConnection(smoother.GetOutputPort())

    if use_flying_edges:
        if using_marching_cubes:
            selector.SetInputArrayToProcess(0, 0, 0, vtkDataObject().FIELD_ASSOCIATION_CELLS,
                                            vtkDataSetAttributes().SCALARS)
        else:
            selector.SetInputArrayToProcess(0, 0, 0, vtkDataObject().FIELD_ASSOCIATION_POINTS,
                                            vtkDataSetAttributes().SCALARS)
    else:
        selector.SetInputArrayToProcess(0, 0, 0, vtkDataObject().FIELD_ASSOCIATION_CELLS,
                                        vtkDataSetAttributes().SCALARS)

    # Strip the scalars from the output
    scalars_off.SetInputConnection(selector.GetOutputPort())
    scalars_off.CopyAttributeOff(vtkMaskFields().POINT_DATA,
                                 vtkDataSetAttributes().SCALARS)
    scalars_off.CopyAttributeOff(vtkMaskFields().CELL_DATA,
                                 vtkDataSetAttributes().SCALARS)

    geometry.SetInputConnection(scalars_off.GetOutputPort())

    writer.SetInputConnection(geometry.GetOutputPort())

    for i in range(start_label, end_label + 1):

        if remove_class is not None:
            if i in remove_class:
                continue

        # see if the label exists, if not skip it
        frequency = histogram.GetOutput().GetPointData().GetScalars().GetTuple1(i)
        if frequency == 0.0:
            continue

        # select the cells for a given label
        selector.SetLowerThreshold(i)
        selector.SetUpperThreshold(i)

        if class_names is not None:
            output_fn = os.path.join(save_root, file_prefix + class_names[i] + '.obj')
        else:
            output_fn = os.path.join(save_root, file_prefix + f'_{i}' + '.obj')

        writer.SetFileName(output_fn)
        writer.Write()

    logging.info(f'Generate models of {file_name} successfully.')


def get_min_max_label(file_paths):
    min_label = 0
    max_label = 0

    files = os.listdir(file_paths)

    for i in files:
        path = os.path.join(file_paths, i)
        label = nib.load(path).get_fdata()
        if min_label > np.min(label):
            min_label = np.min(label)
        if max_label < np.max(label):
            max_label = np.max(label)

    return int(min_label), int(max_label)


def generate(args):
    logging.info('Start generate models.')

    if args.end_label is None:
        _, args.end_label = get_min_max_label(args.file_path)

    path_root = args.file_path
    file_name = os.listdir(path_root)
    file_paths = [os.path.join(path_root, i) for i in file_name]

    if cheak_list_file_format_fit(file_paths):
        splits = file_paths[0].split('.')
        if len(splits) > 1 and splits[2] == 'gz':
            file_format = 'nii.gz'
        else:
            file_format = splits[-1]
        logging.info(f'The file format of the list is {file_format}.')
    else:
        logging.error('The file format of the list is not consistent.')
        return

    save_paths = [os.path.join(args.save_root, i.split('.')[0]) for i in file_name]

    if args.origin is None:
        logging.warning('The origin of the image is not set, the origin will be set to (0, 0, 0).')
        args.origin = (0, 0, 0)
    if args.spacing is None:
        logging.warning('The spacing of the image is not set, the spacing will be set to (1, 1, 1).')
        args.spacing = (1, 1, 1)

    for index, file_path in enumerate(file_paths):
        logging.info(f'Generate {file_path} models.')
        generate_one(file_path, args.start_label, args.end_label, save_paths[index], file_format=file_format,
                     model_format=args.model_format, Origin=args.origin, Spacing=args.spacing,
                     smoothing_iterations=args.smoothing_iterations, pass_band=args.pass_band,
                     feature_angle=args.feature_angle, file_prefix=args.file_prefix,
                     class_names=args.class_names, remove_class=args.remove_class,)

    logging.info('Generate models done.')


if __name__ == '__main__':
    args = parse_args()
    generate(args)
